#!/usr/bin/python
# encoding=utf-8
# cosine 2010/11

"""
Web认证模块：代理设置、HTTP认证、Form认证
测试样例，查看tests目录下的test_t.sina_login.py文件
"""

import time
import urllib2
import urllib
import cookielib
from sgmllib import SGMLParser

class KsspiderError(Exception):
	pass

class Auth():
	"""Web认证类"""
	def __init__(self):
		self.url = None
		self.referer = None
		self.user_agent = None
		self.cookies = None
		
		self.cookiejar = None
		self.cookies_auth = None
		self.proxy_handler = urllib2.BaseHandler()
		self.auth_handler  = urllib2.BaseHandler()
		self.http_opener = None
		
	def build_opener(self):
		"""build urllib2 opener"""
		self.cookiejar = cookielib.LWPCookieJar()
		self.http_opener  = urllib2.build_opener(self.proxy_handler, self.auth_handler, urllib2.HTTPCookieProcessor(self.cookiejar))
		
	def install_opener(self):
		"""install urllib2 opener，谨慎使用，会影响全局urllib2"""
		urllib2.install_opener(self.http_opener)
		
	def set_headers(self,referer='',user_agent='',cookies=''):
		"""设置http请求头"""
		self.referer = referer
		self.user_agent = user_agent
		self.cookies = cookies
		
	def set_http_proxy(self,proxy):
		"""设置http代理 TODO: https"""
		if not proxy:
			self.proxy_handler = urllib2.ProxyHandler({})
			return
		
		__proxy_split = urlparse.urlsplit(proxy)
		__scheme = __proxy_split[0]
		__host_port = __proxy_split[1].split(':')
		__host = __host_port[0]
		__port = None
		if len(__host_port) == 2:
			try:
				__port = int(__host_port[1])
			except:
				pass
		elif len(__host_port) == 1:
			if __scheme == 'http':
				__port = 80
				
		if not __scheme or not __host or not __port:
			err_msg = "proxy value must be in format 'http://url:port'"
			raise KsspiderError,err_msg
			
		__proxy_string = "%s:%d" % (__host, __port)
		self.proxy_handler = urllib2.ProxyHandler({'http':__proxy_string})
		
	def set_http_auth(self,auth_url,auth_type,username,password):
		"""设置http认证：Basic/Digest/NTLM"""
		auth_type = auth_type.lower()
		if auth_type not in ("basic", "digest", "ntlm"):
			err_msg = "KsspiderError: auth type must be in basic/digest/ntlm"
			raise KsspiderError,err_msg
		
		password_mgr =  urllib2.HTTPPasswordMgrWithDefaultRealm()
		password_mgr.add_password(None, auth_url, username, password)
		
		if auth_type == 'basic':
			self.auth_handler = urllib2.HTTPBasicAuthHandler(password_mgr)
		elif auth_type == 'digest':
			self.auth_handler = urllib2.HTTPDigestAuthHandler(password_mgr)
		elif auth_type == 'ntlm':
			try:
				from ntlm import HTTPNtlmAuthHandler
			except:
				err_msg = "requires Python NTLM third-party library in http://code.google.com/p/python-ntlm/"
				raise KsspiderError,err_msg
			
			self.auth_handler = HTTPNtlmAuthHandler.HTTPNtlmAuthHandler(password_mgr)
			
	def set_form_auth(self,auth_url,auth_field_dict):
		"""设置form登录认证"""
		try:
			if not self.http_opener:
				self.build_opener()
			auth_data = urllib.urlencode(auth_field_dict)
			self.http_opener.addheaders = [("Referer",self.referer),("User-Agent",self.user_agent)]
			u=self.http_opener.open(auth_url,auth_data)
			return True
		except Exception,e:
			print 'set_form_auth error: '+str(e)
			return False
		
	def get_cookies_auth(self):
		"""获取认证后的cookies串"""
		set_cookies3_str = self.cookiejar.as_lwp_str()
		set_cookies3_list = set_cookies3_str.split('Set-Cookie3:')
		set_cookies3_list = [cookies3.strip() for cookies3 in set_cookies3_list if cookies3.strip()]
		cookies_str = ''
		for cookies3 in set_cookies3_list:
			cookies_str = cookies_str + cookies3.split(';')[0].replace('"','') + '; '
		#self.cookies_auth = u.headers.dict['set-cookie']
		self.cookies_auth = cookies_str
		return self.cookies_auth
		
	def get_http_opener(self):
		"""获取http请求句柄"""
		if not self.http_opener:
			self.build_opener()
		return self.http_opener
	
	
	
	
